import os
import random
import torchvision.transforms as transforms

from PIL import Image
from torch.utils.data import Dataset

# user-defined package
from networks.utils import utils

class FoodDataset(Dataset):
    def __init__(self, paths):
        super(FoodDataset).__init__()
        self.transforms = None
        self.files = []
        for path in paths:
            self.files += [os.path.join(path,x) for x in os.listdir(path) if x.endswith(".jpg")]
        self.files = sorted(self.files)

    def __len__(self):
        return len(self.files)

    def __getitem__(self,idx):
        fname = self.files[idx]
        im = Image.open(fname)
        try:
            label = int(fname.split("/")[-1].split("_")[0])
        except:
            label = -1 # test has no label

        if self.transforms is not None:
            im = self.transforms(im)
        return im,label


class MixupDataset(Dataset):
    def __init__(self, dataset, alpha=0.8):
        self.dataset = dataset
        self.alpha = alpha

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        # 获取原始图像和标签
        img1, label1 = self.dataset[index]

        # 随机选择另一张图片和标签
        index2 = random.randint(0, len(self.dataset) - 1)
        img2, _ = self.dataset[index2]

        # Mixup图像（线性组合）
        img = self.alpha * img1 + (1 - self.alpha) * img2

        # Mixup标签（线性组合）
        return img, label1

def get_transforms(train_dataset):
    mean, std = utils.calculate_mean_std(train_dataset)

    train_tfm = transforms.Compose([
        transforms.RandomResizedCrop(size=(224, 224), scale=(0.8, 1.0)),
        transforms.RandomHorizontalFlip(0.5),
        transforms.RandomVerticalFlip(0.01),
        transforms.RandomRotation(60),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        transforms.GaussianBlur(kernel_size=2),

        # ToTensor() should be the last one of the transforms
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std) # normalize
    ])

    test_tfm = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std) # normalize
    ])

    return train_tfm, test_tfm

